// -*- mode: c++; indent-tabs-mode: nil; -*-
//
//The Biomolecule Toolkit (BTK) is a C++ library for use in the
//modeling, analysis, and design of biological macromolecules.
//Copyright (C) 2005-2006, Tim Robertson <kid50@users.sourceforge.net>
//
//This program is free software; you can redistribute it and/or modify
//it under the terms of the GNU Lesser General Public License as published
//by the Free Software Foundation; either version 2.1 of the License, or (at
//your option) any later version.
//
//This program is distributed in the hope that it will be useful,  but
//WITHOUT ANY WARRANTY; without even the implied warranty of
//MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
//Lesser General Public License for more details.
//
//You should have received a copy of the GNU Lesser General Public License
//along with this program; if not, write to the Free Software Foundation,
//Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.

#include <iostream>
#include <string>
#include <algorithm>

#include <boost/detail/lightweight_test.hpp>

#include <btk/core/common/exceptions.hpp>
#include <btk/core/molecules/polymer.hpp>
#include <btk/core/molecules/monomer.hpp>
#include <btk/core/atoms/pdb_atom.hpp>
#include <btk/core/io/pdb_system.hpp>

using std::cout;
using std::cerr;
using std::endl;
using std::string;

typedef BTK::ATOMS::PDBAtom atom;
typedef BTK::MOLECULES::Monomer<atom> monomer;
typedef BTK::MOLECULES::Polymer<monomer> polymer;
typedef BTK::IO::PDBSystem<polymer> pdb_system;

using BTK::EXCEPTIONS::BTKException;

void summarize_file(pdb_system const & pdb_file,
                    unsigned & num_chains,
                    unsigned & num_monomers,
                    unsigned & num_atoms)
{
  num_chains = num_monomers = num_atoms = 0;

  for (pdb_system::const_chain_iterator ci = pdb_file.system_begin();
       ci != pdb_file.system_end(); ++ci) {
    ++num_chains;

    for (polymer::const_monomer_iterator mi = ci->polymer_begin();
         mi != ci->polymer_end(); ++mi) {
      ++num_monomers;

      for (monomer::const_atom_iterator ai = mi->structure_begin();
           ai != mi->structure_end(); ++ai) ++num_atoms;
    }
  }
}

int main()
{
  try {

  //
  // Test 1:
  // Load short PDB & make sure that chain, residue
  // and atom loading are correct.
  //
  {
    pdb_system pdb_file("test.short.pdb",
                        false);

    unsigned num_chains = 0;
    unsigned num_monomers = 0;
    unsigned num_atoms = 0;
    std::string header = pdb_file.header();

    BOOST_TEST(std::count(header.begin(),header.end(),'\n') == 2);

    summarize_file(pdb_file,num_chains,num_monomers,num_atoms);

    BOOST_TEST(num_chains == 1);
    BOOST_TEST(num_monomers == 4);
    BOOST_TEST(num_atoms = 30);

    //cout << pdb_file.dictionary() << std::endl;
    cout << "Finshed test 1" << endl;
  }

  //
  // Test 2:
  // Load longer PDB file, and test that multi-chain loading is correct.
  // (w/o HETATM loading)
  //
  {
    pdb_system pdb_file("test.pdb",
                        false);
    
    string header = pdb_file.header();
    unsigned num_chains = 0;
    unsigned num_monomers = 0;
    unsigned num_atoms = 0;

    BOOST_TEST(count(header.begin(),header.end(),'\n') == 2);

    summarize_file(pdb_file,num_chains,num_monomers,num_atoms);

    BOOST_TEST(num_chains == 3);
    BOOST_TEST(num_monomers == 107);
    BOOST_TEST(num_atoms == 1153);

    cout << "Finshed test 2" << endl;
  }

  //
  // Test 3:
  // Same as test 2, with the addition of HETATM loading.
  //
  {
    pdb_system pdb_file("test.pdb",
                        true);

    string header = pdb_file.header();
    unsigned num_chains = 0;
    unsigned num_monomers = 0;
    unsigned num_atoms = 0;

    BOOST_TEST(count(header.begin(),header.end(),'\n') == 2);

    summarize_file(pdb_file,num_chains,num_monomers,num_atoms);

    BOOST_TEST(num_chains == 5);
    BOOST_TEST(num_monomers == (107 + 258));
    BOOST_TEST(num_atoms == (1153 + 258));

    cout << "Finshed test 3" << endl;
  }

  //
  // Test 4:
  // Multiple models, w/o HETATMs.
  //
  {
    pdb_system pdb_file("multi_model_test.pdb",
                        false);

    unsigned num_models = 1;
    unsigned num_chains = 0;
    unsigned num_monomers = 0;
    unsigned num_atoms = 0;

    while (pdb_file.get_next_model(false)) {
      summarize_file(pdb_file,num_chains,num_monomers,num_atoms);
      ++num_models;

      BOOST_TEST(num_chains == 1);
      BOOST_TEST(num_monomers == 77);
      BOOST_TEST(num_atoms = 1194);

      // Do a few spot-checks, just to make sure that the parsing is
      // working correctly for each model.
      if (num_models == 2) {
        monomer & first_monomer =
          *(pdb_file.system_begin()->polymer_begin());
        atom & first_atom = *(first_monomer.structure_begin());
 
        BOOST_TEST(first_monomer.name() == "LEU");
        BOOST_TEST(first_atom.name() == "N");
        BOOST_TEST(first_atom.position()[0] == 4.716);
      } else if (num_models == 5) {
        monomer & first_monomer =
          *(pdb_file.system_begin()->polymer_begin());
        atom & second_atom = *(++first_monomer.structure_begin());

        BOOST_TEST(first_monomer.name() == "LEU");
        BOOST_TEST(second_atom.name() == "CA");
        BOOST_TEST(second_atom.position()[2] == -5.629);
      } else if (num_models == 10) {
        monomer & last_monomer =
          *(--pdb_file.system_begin()->polymer_end());
        atom & second_atom = *(++last_monomer.structure_begin());

        BOOST_TEST(last_monomer.name() == "ASN");
        BOOST_TEST(second_atom.name() == "CA");
        BOOST_TEST(second_atom.position()[1] == 11.927);
      }
    }

    BOOST_TEST(num_models == 10);

    cout << "Finshed test 4" << endl;
  }

  //
  // Test 5:
  // The whole enchilada.  Multiple models w/het atoms.
  //
  {
    pdb_system pdb_file("multi_model_test.pdb",
                        true);

    unsigned num_models = 1;
    unsigned num_chains = 0;
    unsigned num_monomers = 0;
    unsigned num_atoms = 0;

    while (pdb_file.get_next_model()) {
      summarize_file(pdb_file,num_chains,num_monomers,num_atoms);
      ++num_models;

      BOOST_TEST(num_chains == 2);
      BOOST_TEST(num_monomers == 80);
      BOOST_TEST(num_atoms == 1269);
    }

    cout << "Finshed test 5" << endl;
  }

  } catch (BTKException & e) {
    cerr << "Caught BTKException: " << e.what() << endl;
  }

  return boost::report_errors();
};

